import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import matplotlib
matplotlib.use('Agg')
import numpy as np
np.random.seed(42)

from configs.config import load_config
from data_process.structure import RtkData
from data_process.tensor_io import TensorIO
from data_process.rtk_reader import RtkReader
from data_process.plot_bev import BEVGenerator
from data_process.radar_reader import RadarReader
from data_process.frame_selector import FrameSelector
from data_process.plot_3d_points import PointCloudVisualizer
from utils.coordinate_transformer import CoordinateTransformer
from data_process.radar_to_tensor_converter import RadarToTensorConverter


if __name__ == "__main__":
    cfg = load_config("./configs/data_process.py")
    reader = RadarReader(cfg)
    rtk_reader = RtkReader(cfg)

    for data_name, folder_path in cfg.get('data_root_list').items():
        folder_path = os.path.join(folder_path, 'SARRadar')
        with os.scandir(folder_path) as it:
            bin_file_count = sum(1 for entry in it if entry.is_file() and entry.name.endswith('.bin'))
        batch = 0
        while bin_file_count > 40:
            bin_file_count -= cfg.get('data_batch_size')
            batch += 1

            if batch <= 39 and data_name == 'data331':
                continue
            data_name_with_batch = f'{data_name}_b{batch}'

            # 读取数据
            cfg['data_name'] = data_name_with_batch
            radar_data = reader.process_all_files(folder_path, batch) # z [-0.3, 8], distance [6.5, 85], c [0.14, 1]
            rtk_data = rtk_reader.process_all_files(folder_path, radar_data.keys())

            radar_data = dict(sorted(radar_data.items(), key=lambda item: int(item[0].split('_')[1])))
            rtk_data = dict(sorted(rtk_data.items(), key=lambda item: int(item[0].split('_')[1])))

            if False: # and not os.path.exists(os.path.join(cfg.get('bev_output_folder'), f'total_BEV_{cfg.get("data_name")}.png')):
                # 画大图
                ref_rtk_id = min(rtk_data.keys(), key=lambda x: int(x.split('_')[1]))
                ref_rtk = rtk_data[ref_rtk_id]
                ref_rtk = RtkData(
                    latitude=ref_rtk.latitude,
                    longitude=ref_rtk.longitude,
                    height=ref_rtk.height,
                    heading=0.0,
                    roll=0.0,
                    pitch=0.0,
                    timeStamp=0,
                )
                
                # radar_data = FrameSelector.select_frame_and_points(radar_data, rtk_data)
                radar_data = CoordinateTransformer.transform_points(radar_data, rtk_data, ref_rtk)
                cfg['lat0'] = ref_rtk.latitude
                cfg['lon0'] = ref_rtk.longitude
                cfg['h0'] = ref_rtk.height
                cfg['heading'] = ref_rtk.heading

                PointCloudVisualizer.visualize_3d_points(radar_data, cfg, id=f'3d_{cfg.get("data_name")}')
                BEVGenerator.create_bev(
                    radar_data,
                    cfg,
                    id=f'total_BEV_{cfg.get("data_name")}',
                    images=['all', 'z02', 'z00.5', 'z-10.5', 's20w', 's40w'],
                    output_folder=cfg.get('bev_output_folder'),
                    save_meta=True)

            if True:
                # 准备tensor
                sequences = FrameSelector.select_sequences(rtk_data, cfg)
                tensors = RadarToTensorConverter.convert(radar_data, rtk_data, sequences, cfg)
                TensorIO.save_tensors(tensors, cfg)
